#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>
#include <ngx_http_cache_control_filter_module.h>
#include <db.h>

typedef struct {
    ngx_str_t db_env_path;
    ngx_str_t db_path;
    DB_ENV *db_env;
    DB *db;
} bdb_ctx_t;


static ngx_int_t
ngx_http_bdb_init_module(ngx_cycle_t *c, ngx_http_cache_control_filter_source_t *s)
{
    return NGX_OK;
}

static ngx_int_t
bdb_open_env(ngx_log_t *log, ngx_http_cache_control_filter_source_t *s)
{
    ngx_int_t ret;
    DB_ENV *db_envp;
    DB *dbp;
    bdb_ctx_t *ctx;
    ngx_uint_t open_flags;

    ctx = s->priv_data;

    /* create db enviroment handler */
    ret = db_env_create(&db_envp, 0);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, log, 0,
                      "Create database enviroment handler failed. %s",
                      db_strerror(ret));
        return NGX_ERROR;
    }

    (void)db_envp->set_cachesize(db_envp, 0, s->cache_size * 1024, 0);
    (void)db_envp->set_flags(db_envp, DB_TXN_NOSYNC, 1);

    /* open the enviroment */
    open_flags = DB_CREATE |    /* Create the environment if it does
                                 * not already exist. */
        DB_INIT_TXN  | /* Initialize transactions */
        DB_INIT_LOCK | /* Initialize locking. */
        DB_INIT_LOG  | /* Initialize logging */
        DB_INIT_MPOOL| /* Initialize the in-memory cache. */
        DB_THREAD    |
        DB_INIT_REP  |
        DB_RECOVER;
    ret = db_envp->open(db_envp, (const char *)ctx->db_env_path.data, open_flags, 0);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, log, 0,
                      "Open database enviroment handler failed for path '%s'. %s",
                      ctx->db_env_path.data, db_strerror(ret));
        return NGX_ERROR;
    }

    /* create db handler */
    ret = db_create(&dbp, db_envp, 0);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, log, 0,
                          "Create database handler failed. %s",
                      db_strerror(ret));
        return NGX_ERROR;
    }

    open_flags = DB_CREATE | DB_AUTO_COMMIT;
    ret = dbp->open(dbp,
                    NULL,
                    (char *)ctx->db_path.data,
                    NULL,
                    DB_BTREE,
                    open_flags,
                    0);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, log, 0,
                      "Open database '%s' failed. %s",
                      ctx->db_path.data, db_strerror(ret));
        return NGX_ERROR;
    }

    ctx->db_env = db_envp;
    ctx->db = dbp;

    return ret;
}

static ngx_int_t
bdb_close_env(ngx_log_t *log, ngx_http_cache_control_filter_source_t *s)
{
    ngx_int_t ret = NGX_OK;
    DB_ENV *db_envp;
    DB *dbp;
    bdb_ctx_t *ctx;

    ctx = (bdb_ctx_t *)s->priv_data;

    db_envp = ctx->db_env;
    dbp = ctx->db;

    if (dbp != NULL) {
        ret = dbp->close(dbp, 0);
        if (ret) {
            ngx_log_error(NGX_LOG_ERR, log, 0,
                          "Database '%s' close failed. %s",
                          ctx->db_path.data, db_strerror(ret));
        }
    }

    if (db_envp != NULL) {
        ret = db_envp->close(db_envp, 0);
        if (ret) {
            ngx_log_error(NGX_LOG_ERR, log, 0,
                          "Database enviroment '%s' close failed. %s",
                          ctx->db_env_path.data, db_strerror(ret));
        }
    }

    return ret;
}

static ngx_int_t
bdb_reopen_db(ngx_log_t *log, ngx_http_cache_control_filter_source_t *s)
{
    ngx_int_t ret = NGX_OK;
    ngx_uint_t open_flags;
    DB *dbp;
    DB_ENV *db_envp;
    bdb_ctx_t *ctx;

    ctx = (bdb_ctx_t *)s->priv_data;

    db_envp = ctx->db_env;
    dbp = ctx->db;

    if (dbp != NULL) {
        ret = dbp->close(dbp, 0);
        if (ret) {
            ngx_log_error(NGX_LOG_ERR, log, 0,
                          "Database '%s' close failed. %s",
                          ctx->db_path.data, db_strerror(ret));
            return NGX_ERROR;
        }
    }

    /* create db handler */
    ret = db_create(&dbp, db_envp, 0);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, log, 0,
                          "Create database handler failed. %s",
                      db_strerror(ret));
        return NGX_ERROR;
    }

    open_flags = DB_CREATE | DB_AUTO_COMMIT;
    ret = dbp->open(dbp,
                    NULL,
                    (char *)ctx->db_path.data,
                    NULL,
                    DB_BTREE,
                    open_flags,
                    0);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, log, 0,
                      "Open database '%s' failed. %s",
                      ctx->db_path.data, db_strerror(ret));
        return NGX_ERROR;
    }

    ctx->db = dbp;

    return ret;
}

static ngx_int_t
ngx_http_bdb_init_process(ngx_cycle_t *c, ngx_http_cache_control_filter_source_t *s)
{
    ngx_int_t ret;
    bdb_ctx_t *ctx;

    if (s->priv_data == NULL) {
        s->priv_data = ngx_palloc(c->pool, sizeof(bdb_ctx_t));
        if(s->priv_data == NULL) {
            return NGX_ERROR;
        }
    }

    if (!ngx_strcmp(s->path.data, "")) {
        return NGX_OK;
    }

    ctx = (bdb_ctx_t *)s->priv_data;

    /* generate the paths */
    ctx->db_env_path.data = ngx_palloc(c->pool, s->path.len + 1);
    if (ctx->db_env_path.data == NULL) {
        return NGX_ERROR;
    }
    ngx_cpystrn(ctx->db_env_path.data, s->path.data, s->path.len + 1);
    ctx->db_env_path.len = s->path.len;

    ctx->db_path.data = (u_char *)"pair.db";
    ctx->db_path.len = strlen("pair.db");

    ret = bdb_open_env(c->log, s);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, c->log, 0,
                      "Open bdb enviroment failed.");
        return ret;
    }

    return NGX_OK;
}

static ngx_int_t
ngx_http_bdb_exit_process(ngx_cycle_t *c, ngx_http_cache_control_filter_source_t *s)
{
    ngx_int_t ret = NGX_OK;
    if (s->priv_data == NULL) {
        return NGX_OK;
    }

    ret = bdb_close_env(c->log, s);
    if (ret) {
        ngx_log_error(NGX_LOG_ALERT, c->log, 0,
                      "Close bdb enviroment failed.");
        return ret;
    }

    return ret;
}

static ngx_int_t
ngx_http_bdb_get_expire_pair(ngx_http_request_t *r, ngx_http_cache_control_filter_source_t *s, time_t *e)
{
    ngx_int_t ret, err = NGX_OK;
    DB *dbp;
    DBT key, value;
    u_char *iter;
    ngx_uint_t cache_time;
    ngx_connection_t *c;

    if (s->priv_data == NULL) {
        return NGX_OK;
    }

    c = r->connection;
    dbp = ((bdb_ctx_t *)s->priv_data)->db;

    if (dbp == NULL) {
        *e = 0;
        return NGX_OK;
    }

    memset(&key, 0, sizeof(DBT));
    memset(&value, 0, sizeof(DBT));
    key.data = r->uri.data;
    value.data = &cache_time;
    value.ulen = sizeof(cache_time);
    value.flags = DB_DBT_USERMEM;

    key.size = r->uri.len;
    iter = r->uri.data + r->uri.len - 1;
    if (s->match_type == NGX_HTTP_CACHE_CONTROL_FILTER_MATCH_TYPE_PREFIX) {
        for(;;) {
            ret = dbp->get(dbp, NULL, &key, &value, 0);
            /* if we use replication in our enviroment, after rollback, the original hanlder will be invalid
               and, we should reopen a new one to operate */
            if (ret == DB_REP_HANDLE_DEAD) {
                ret = bdb_reopen_db(c->log, s);
                if (ret) {
                    ngx_log_error(NGX_LOG_ERR, c->log, 0,
                                  "Reopen database failed.");
                    return NGX_ERROR;
                }

                ngx_log_error(NGX_LOG_INFO, c->log, 0,
                              "Reopen database successfully.");
                dbp = ((bdb_ctx_t *)s->priv_data)->db;
                continue;
            } else if (ret == DB_NOTFOUND) {
                /* iter reach the start of the uri */
                if (iter == r->uri.data) {
                    ngx_log_error(NGX_LOG_INFO, c->log, 0,
                              "No corresponding cache time for"
                                  " uri '%s'.", r->uri.data);
                    err = NGX_ERROR;
                    break;
                }
                /* find the prefix split by the '/' */
                for(;;) {
                    iter--;
                    key.size--;
                    if (*iter == '/' || iter == r->uri.data) {
                    break;
                    }
                }
            } else if (ret) {
                ngx_log_error(NGX_LOG_ALERT, c->log, 0,
                              "Get cache time for %s failed. %s",
                              r->uri.data, db_strerror(ret));
                break;
                err = NGX_ERROR;
            } else {
                *e = cache_time;
                err = NGX_OK;
                break;
            }
        }
    } else {
    REOPEN:
        ret = dbp->get(dbp, NULL, &key, &value, 0);
        if (ret == DB_REP_HANDLE_DEAD) {
            ret = bdb_reopen_db(c->log, s);
            if (ret) {
                ngx_log_error(NGX_LOG_ERR, c->log, 0,
                              "Reopen database failed.");
                return NGX_ERROR;
            }

            ngx_log_error(NGX_LOG_INFO, c->log, 0,
                          "Reopen database successfully.");
            dbp = ((bdb_ctx_t *)s->priv_data)->db;
            goto REOPEN;
        } else if(ret == DB_NOTFOUND) {
            ngx_log_error(NGX_LOG_INFO, c->log, 0,
                          "No corresponding cache time for"
                          " uri '%s'.", r->uri.data);
            err = NGX_ERROR;
        } else if (ret) {
            ngx_log_error(NGX_LOG_ALERT, c->log, 0,
                          "Get cache time for %s failed. %s",
                          r->uri.data, db_strerror(ret));
            err = NGX_ERROR;
        } else {
            *e = cache_time;
            err = NGX_OK;
        }
    }

    return err;
}

ngx_http_cache_control_filter_supported_source_t supported_sources[] = {
    {"bdb",
     {
         ngx_http_bdb_init_module,
         ngx_http_bdb_init_process,
         ngx_http_bdb_exit_process,
         ngx_http_bdb_get_expire_pair
     }
    },
    {NULL, {NULL, NULL, NULL, NULL}}
};
